import cv2
import numpy as np
import os
from torch.utils import data as data
from os import path as osp

from basicsr.data.transforms import augment, paired_random_crop
from basicsr.utils import FileClient, img2tensor
from basicsr.utils.registry import DATASET_REGISTRY


@DATASET_REGISTRY.register()
class DehazingDataset(data.Dataset):
    """
    A Dataset class for image dehazing, designed for the basicsr framework.
    It pairs a hazy image (LQ) with its corresponding clear ground-truth image (GT).
    The pairing is based on the filename convention, e.g., '0001_0.8_0.2.png' in the LQ folder
    corresponds to '0001.png' in the GT folder.

    Args:
        opt (dict): Config for training datasets. It contains the following keys:
            dataroot_gt (str): Data root path for ground-truth (clear) images.
            dataroot_lq (str): Data root path for low-quality (hazy) images.
            gt_size (int): Cropped patched size for gt patches.
            use_flip (bool): Use horizontal flips.
            use_rot (bool): Use rotation.
    """

    def __init__(self, opt):
        super(DehazingDataset, self).__init__()
        self.opt = opt
        # File client (for reading data from different backends, e.g., disk, lmdb, ceph)
        self.file_client = None
        self.io_backend_opt = opt['io_backend']

        # Paths for LQ (hazy) and GT (clear) images
        self.lq_folder = opt['dataroot_lq']
        self.gt_folder = opt['dataroot_gt']
        
        # Scan the LQ folder to get the list of hazy image paths
        # We assume that for every LQ image, there's a corresponding GT image.
        self.lq_paths = sorted([osp.join(self.lq_folder, f) for f in os.listdir(self.lq_folder)])

    def __getitem__(self, index):
        # -------------------------------- Read LQ and GT images -------------------------------- #
        # Get the path of the LQ (hazy) image
        lq_path = self.lq_paths[index]
        
        # Derive the corresponding GT image path based on the filename convention
        lq_filename = osp.basename(lq_path)
        gt_filename = lq_filename.split('_')[0] + '.jpg'  # Assumes GT is always .png
        gt_path = osp.join(self.gt_folder, gt_filename)

        # Read images using OpenCV
        img_lq = cv2.imread(lq_path, cv2.IMREAD_COLOR).astype(np.float32) / 255.
        img_gt = cv2.imread(gt_path, cv2.IMREAD_COLOR).astype(np.float32) / 255.
        
        # Convert BGR images read by cv2 to RGB
        img_lq = img_lq[:, :, [2, 1, 0]]
        img_gt = img_gt[:, :, [2, 1, 0]]

        # -------------------------------- Data Augmentation -------------------------------- #
        # For image restoration tasks like dehazing, the scale is 1
        scale = self.opt.get('scale', 1)
        gt_size = self.opt['gt_size'] 

        # Paired random crop
        img_gt, img_lq = paired_random_crop(img_gt, img_lq, gt_size, scale, gt_path)
        
        # Paired random flip and rotation
        img_gt, img_lq = augment([img_gt, img_lq], self.opt['use_flip'], self.opt['use_rot'])
        
        # -------------------------------- Convert to Tensor -------------------------------- #
        # Convert numpy arrays to tensors
        # bgr2rgb is False because we have already converted from BGR to RGB.
        img_gt, img_lq = img2tensor([img_gt, img_lq], bgr2rgb=False, float32=True)
        
        return {
            'lq': img_lq,
            'gt': img_gt,
            'lq_path': lq_path,
            'gt_path': gt_path
        }

    def __len__(self):
        return len(self.lq_paths)

